from action import *
from domainparser import *
import parserbasic
import time
import re

#regular expressions for reading preconditions
jacobianRE = re.compile('J\((.+)\)\s*=\s*(.+)')
tolRE = re.compile('TOL\s*=\s*(.+)')

class ActionParser(object):
    """A class that parses actions from a file.  Instantiated with the global
    environment, and a state description"""
    def __init__(self,env,stateNames):
        self.env = env
        self.stateNames = stateNames
        pass
        
    def parse(self, actionName, code, start):
        """Returns a new Action defined by the given code"""
        linenos,sections = parserbasic.split_sections(code,['PARAMS:','PRECOND:','EFFECTS:','COST:'],start+1)
        if len(linenos)==0:
            raise RuntimeError("STRIPS: Action needs params, preconditions, "\
                               "and effects:\n" + ''.join(code))

        a = Action(actionName)
        if 'PARAMS:' in linenos:
            self.parse_params(a, sections['PARAMS:'], linenos['PARAMS:'][0])

        if 'PRECOND:' in linenos:
            self.parse_preconds(a, sections['PRECOND:'], linenos['PRECOND:'][0])

        if 'EFFECTS:' in linenos:
            self.parse_effects(a, sections['EFFECTS:'], linenos['EFFECTS:'][0])

        if 'COST:' in linenos:
            self.parse_costs(a,sections['COST:'],linenos['COST:'][0])
                               
        return a
            
    def parse_params(self, a, code, start):
        for line in code:
            #check for empty line
            if len(line.strip())==0: continue
            
            lineParts = [l.strip() for l in line.split(' in ',1)]
            if len(lineParts) != 2:
                raise RuntimeError("STRIPS: params line '"+line+"' is not formatted according to VAR in DOMAIN")
            varName, varVal = lineParts
            if varName in a.params:
                raise RuntimeError("STRIPS: Parameter name collision: " + varName)
            if varName in self.env:
                print "Warning, global namespace collision:",varName
                print "Value in global namespace:",self.env[varName]
            tempVars = set(self.stateNames.copy())
            tempVars |= set(a.params.keys())
            a.addParam(varName,DomainParser.parse(self.env, varVal,tempVars))

    def parse_precond(self,line,paramNames):
        if line.startswith('EQ'):           
            line = line[2:]
            lineParts = [l.strip() for l in line.split(';')]
            functionParts = [l.strip() for l in lineParts[0].split('==')]
            if len(functionParts) != 2:
                raise RuntimeError("Expected exactly one equality comparison")
            fnCode,val = functionParts
            if not self.is_zero(self.env, val):
                raise RuntimeError("Vector field function must target a zero vector")
            tol = 1e-4
            jacobians = {}
            for option in lineParts[1:]:
                if tolRE.match(option):
                    m = tolRE.match(option)
                    tol = float(m.group(1))
                elif jacobianRE.match(option):
                    m = jacobianRE.match(option)
                    paramName = m.group(1)
                    expr = m.group(2)
                    if paramName not in paramNames:
                        raise RuntimeError("Jacobian is not measured with respect to an action parameter")
                    jacobians[paramName] = expr
                else:
                    print option
                    raise RuntimeError("Invalid suffix on equality precondition")
            pre = Precondition()
            pre.setEquality(fnCode,self.env,self.stateNames,paramNames,tol,jacobians)
            return pre
        elif line.startswith('MP'):
            line = line[2:]
            pre = Precondition()
            pre.setMotionPlanning(line)
            return pre
        else:
            pre = Precondition()
            pre.setCode(line,self.env,self.stateNames,paramNames)
            return pre
                
    def parse_preconds(self, a, code, start):
        paramNames = [p for p in a.params]
        for line in code:           
            line = line.strip()
            #check for empty line
            if len(line)==0: continue
            a.addPrecond(self.parse_precond(line,paramNames))
            
    def parse_effects(self, a, code, start):
        for line in code:
            #check for empty line
            if len(line.strip())==0: continue

            lineParts = [l.strip() for l in line.split('=')]
            if len(lineParts) != 2:
                raise RuntimeError("STRIPS: effects line '"+line+"' is not formatted according to VAR = VALUE")
            varName, varVal = lineParts
            # Kind of a hack to see if we're assigning to a subscript or
            #  dot member
            varName = varName.split('[')[0]
            varName = varName.split('.')[0]
            if varName not in self.stateNames:
                raise RuntimeError("Effects must assign to state variables:\n" + line)
            a.addEffect(line.strip())

    def parse_costs(self, a, code, start):
        for line in code:
            #check for empty line
            if len(line.strip())==0: continue
            try:
                c = eval(line)
                if isinstance(c,(int,float)):
                    a.costTerms.append(c)
                else:
                    raise NameError("Blah")
            except NameError:
                a.costTerms.append(line)
            
    def is_zero(self, env, val):
        if isinstance(val, int) or isinstance(val, float):
            return val == 0
        if isinstance(val, list) or isinstance(val, tuple):
            for v in val:
                if not self.is_zero(env, v):
                    return False
            return True
        # Don't allow delayed evaluations here
        if not isinstance(val, str):
            raise RuntimeError("Wrong variable type: " + type(val))
        return self.is_zero(env, eval(val, globals(), env))

def testActionSampler(action,env,state,n=1000):
    print
    print "*** Testing ActionParameterSampler ***"
    
    asamp = ActionParameterSampler(action,env)
    asamp.bindState(state)
    t0 = time.time()
    numfeas = 0
    for i in xrange(n):
        params = asamp.sample_domains()
        if action.applicable(env,params,state):
            numfeas += 1
    print "% feasible raw sampled params:",float(numfeas)/n,", time",time.time()-t0

    t0 = time.time()
    numfeas = 0
    for i in xrange(n):
        params = asamp.sample_applicable(10)
        if params != None:
            if action.applicable(env,params,state):
                numfeas += 1
    print "% feasible solved/sampled params:",float(numfeas)/n,", time",time.time()-t0
    

def testAction(parser,code,state,goodparams=None,badparams=None,n=1000):
    print
    print "*** Testing ActionParser ***"
    a = parser.parse('test',code,0)
    print "Name:",a.name
    print "Params:"
    for p in a.params:
        print "\t",p,"in",a.params[p]
    print "Preconds:"
    for p in a.preconds:
        print "\t",p
    print "Effects:"
    for e in a.effects:
        print "\t",e

    if goodparams:
        for p in goodparams:
            if a.applicable(parser.env,p,state)!=True:
                print "Good parameters",p,"not applicable to state",state
                for pre in a.preconds:
                    if not pre.applicable(parser.env,p,state):
                        print pre,'not satisfied'
                for (pname,val) in p.items():
                    if not a.params[pname].contains(val):
                        print pname,'=',val,'not in',a.params[pname]
                raise ValueError("Good parameter value not applicable")
            print 'Result of',BoundAction(a,p)
            print "\t",a.apply(parser.env,p,state)
    if badparams:
        for p in badparams:
            if a.applicable(parser.env,p,state)!=False:
                print "Bad parameters",p,"applicable to state",state
                raise ValueError("Bad parameter value was applicable")

    testActionSampler(a,parser.env,state,n=n)

#debug the code
if __name__ == '__main__':
    env = globals()
    stateNames = set(['q','index'])
    parser = ActionParser(env,stateNames)

    #construct and test an action that:
    #-allows q[1] to be set to something else and increments index.
    #-requires q[0] to stay the same
    code = \
"""PARAMS:
    qdes in [(0,1),(0,1)]
    inew in {1,2,3}
PRECOND:
    EQ qdes[0]-q[0]==0
    inew == index+1
EFFECTS:
    q[1] = qdes[1]
    index = inew
"""
    state = {'q':[0.3,0.1],'index':2}
    goodparams = [{'qdes':[0.3,0.5],'inew':3},{'qdes':[0.3,0.9],'inew':3}]
    badparams = [{'qdes':[0.2,0.5],'inew':3},{'qdes':[0.3,0.5],'inew':2}]
    testAction(parser,code,state,goodparams,badparams)

    #construct and test an action that:
    #-has an equality condition that depends on a discrete value
    code = \
"""PARAMS:
    qdes in [(0,5),(0,1)]
    inew in {1,2,3}
PRECOND:
    EQ qdes[0]-inew-q[0]==0
EFFECTS:
    q = qdes
    index = inew
"""
    state = {'q':[0.3,0.1],'index':2}
    goodparams = [{'qdes':[1.3,0.5],'inew':1},{'qdes':[2.3,0.9],'inew':2}]
    badparams = [{'qdes':[0.3,0.5],'inew':0},{'qdes':[1.3,1.1],'inew':1}]
    testAction(parser,code,state,goodparams,badparams)
